#include "track.h"
#include <iostream>
#include <algorithm>
#include <numeric>
#include <set>

CentroidTracking::CentroidTracking(int maxDisappeared, double distanceThreshold) :
	maxDisappeared(maxDisappeared), nextObjectId(0), distanceThreshold(distanceThreshold) {}

void CentroidTracking::registerObject(const cv::Point& centroid, const cv::Rect& rect) {
	objects[nextObjectId] = centroid;
	disappeared[nextObjectId] = 0;
	frameCount[nextObjectId] = 1;
	dxdy[nextObjectId] = cv::Point(0, 0);
	outrects[nextObjectId] = rect;

	nextObjectId++;
}

void CentroidTracking::dergisterObject(int objectID) {
	objects.erase(objectID);
	disappeared.erase(objectID);
	frameCount.erase(objectID);
	dxdy.erase(objectID);
	outrects.erase(objectID);
}

TrackingResult CentroidTracking::update(const std::vector<cv::Rect>& rects) {
	TrackingResult result;

	if (rects.empty()) {
		for (const auto& objectID : disappeared) {
			disappeared[objectID.first]++;
			frameCount[objectID.first] = 0;
			dxdy[objectID.first] = cv::Point(0, 0);

			if (disappeared[objectID.first] > maxDisappeared) {
				dergisterObject(objectID.first);
			}
		}
		
		result.objects = objects;
		result.disappeared = disappeared;
		result.frameCount = frameCount;
		result.outrects = outrects;

		return result;
	}

	std::vector<cv::Point> inputCentroids;
	std::vector<cv::Rect> inputRects;

	for (const cv::Rect& rect : rects) {
		cv::Point centroid(rect.x + rect.width / 2, rect.y + rect.height / 2);
		inputCentroids.push_back(centroid);
		inputRects.push_back(rect);
	}

	if (objects.empty()) {
		for (size_t i = 0; i < inputCentroids.size(); i++) {
			registerObject(inputCentroids[i], inputRects[i]);
		}
	}
	else {
		std::vector<int> objectIDs;
		std::vector<cv::Point> objectCentroids;

		for (const auto& object : objects) {
			objectIDs.push_back(object.first);
			objectCentroids.push_back(object.second);
		}

		std::vector<cv::Point> dxdys;
		dxdys.reserve(objectCentroids.size());

		for (size_t i = 0; i < objectCentroids.size(); i++) {
			const cv::Point& obv = objectCentroids[i];
			const cv::Point& dxdyValue = dxdy[objectIDs[i]];
			dxdys.push_back(obv + dxdyValue);
		}

		objectCentroids = dxdys;

		cv::Mat D(objectCentroids.size(), inputCentroids.size(), CV_64F);

		for (size_t i = 0; i < objectCentroids.size(); i++) {
			for (size_t j = 0; j < inputCentroids.size(); j++) {
				double distance = cv::norm(objectCentroids[i] - inputCentroids[j]);
				D.at<double>(i, j) = distance;
			}
		}

		std::vector<int> rows(D.rows);
		std::iota(rows.begin(), rows.end(), 0);

		std::vector<int> cols(D.cols);
		std::iota(cols.begin(), cols.end(), 0);

		std::set<int> usedRows;
		std::set<int> usedCols;

		for (int row : rows) {
			for (int col : cols) {
				if (usedRows.count(row) || usedCols.count(col)) {
					continue;
				}

				if (D.at<double>(row, col) > distanceThreshold) {
					continue;
				}

				int objectID = objectIDs[row];
				dxdy[objectID] = inputCentroids[col] - objects[objectID];
				objects[objectID] = inputCentroids[col];
				outrects[objectID] = inputRects[col];
				disappeared[objectID] = 0;
				frameCount[objectID]++;

				usedRows.insert(row);
				usedCols.insert(col);
			}
		}

		std::set<int> unusedRows;
		std::set<int> unusedCols;

		for (int i = 0; i < D.rows; i++) {
			if (usedRows.count(i) == 0) {
				unusedRows.insert(i);
			}
		}

		for (int i = 0; i < D.cols; i++) {
			if (usedCols.count(i) == 0) {
				unusedCols.insert(i);
			}
		}

		if (D.rows >= D.cols) {
			for (int row : unusedRows) {
				int objectID = objectIDs[row];
				disappeared[objectID]++;
				frameCount[objectID] = 0;
				dxdy[objectID] = cv::Point(0, 0);

				if (disappeared[objectID] > maxDisappeared) {
					dergisterObject(objectID);
				}
			}
		}
		else {
			for (int col : unusedCols) {
				registerObject(inputCentroids[col], inputRects[col]);
			}
		}
	}

	for (const auto& object : objects) {
		result.objects[object.first] = object.second;
	}

	for (const auto& item : frameCount) {
		result.frameCount[item.first] = item.second;
	}

	for (const auto& item : outrects) {
		result.outrects[item.first] = item.second;
	}

	return result;
}